from src.module.Model import Model
from src.module.Embedding import Embedding

from numpy import dot
from numpy.linalg import norm

from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores.faiss import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain

from src.constance import SEARCH_ARTICLE_TEXT


## 问题内容余弦相似度示例
class SimilarityMatch:
    def __init__(self):
        self.embedding = Embedding().embedding

    ## 将数据内容转换为向量
    def generator_data(self, value: str):
        return self.embedding.embed_documents([value])[0]

    ## 计算余弦相似度
    @staticmethod
    def cosine_similarity(vec1, vec2):
        return dot(vec1, vec2) / (norm(vec1) * norm(vec2))

    ## 进行计算
    def start(self):
        content = [ '小猫是可爱的', '小狗是忠诚的', '小老鼠是比较小的' ]
        query = '可爱的'

        # 将内容转换为向量
        document_embedding = [ self.generator_data(item) for item in content ]
        query_embedding = self.generator_data(query)

        ## 将 query 问题内容与数据内容进行转换，得出每一个 content 的相似度比分
        num_list = [self.cosine_similarity(query_embedding, item) for item in document_embedding]

        ## 拿到数组中最大那个值的 index
        max_rate = max(num_list)
        max_index = num_list.index(max_rate)

        # 余弦相似度得出结论
        ## 给定相似度阈值，如果相似度过小，表明文本不存在知识库中
        if max_rate < 0.6:
            print('未能找到相似的文本内容')
        else:
            print('content:', content[max_index], 'rate is:', max_rate)

## 意图匹配示例
class IntentMatch:
    def __init__(self):
        self.embedding = Embedding().embedding
        self.model = Model().model

    def parser_file(self):
        docs = TextLoader(SEARCH_ARTICLE_TEXT, encoding='utf8').load()
        text_splitter  = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20).split_documents(docs)
        retriever = FAISS.from_documents(text_splitter, self.embedding)
        return retriever

    @staticmethod
    def get_prompt():
        return ChatPromptTemplate.from_template("""
            请根据提供的上下文回复问题：
            <context>{context}<context>
            
            Question: {input}
        """)

    def start(self):
        retriever = self.parser_file().as_retriever() # 向量数据库句柄
        prompt = self.get_prompt() # 提示词

        # 文档 chain, 上下文的内容固定为 context 标签以及 context 文本内容
        dos_chain = create_stuff_documents_chain(self.model, prompt)

        # Rag 专用chain，一般配合 create_stuff_documents_chain 一起使用
        # 缺点：只能单轮回答，不支持多轮回答，因此使用场景一般在检索链。
        # 性能相比： ConversationalRetrievalChain 较好，但是一般的使用都还是使用 ConversationalRetrievalChain
        retrieval_chain = create_retrieval_chain(retriever, dos_chain)

        response = retrieval_chain.invoke({ "input": "请给我总结一下，第一个故事" })
        print(response)


class EmbeddingStorage:
    def __init__(self):
        # self.similarityExample = SimilarityMatch() ## 余弦相似度
        self.intentExample =  IntentMatch()  ## 意图匹配示例

    def start(self):
        # self.similarityExample.start()
        self.intentExample.start()